from contextlib import contextmanager
from enum import Enum
import os
import re
import tempfile
import time
from typing import List
from common.rpc.paperfile import paperfile_pb2
import fitz
from minio.error import S3Error

from app.jobs.pdfparser.internal.logic.splitter.text_splitter import TextSegment, TextSplitter
from app.jobs.pdfparser.internal.logic.task import KqPaperMeta, PDFTask, PDFTaskType
from app.jobs.pdfparser.internal.svc.servicecontext import ServiceContext
from common.utils import logger
import pymupdf4llm

class ParseStatus(Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    SUCCESS = "success"
    FAILED = "failed"

class PDFParseTask(PDFTask):
    def __init__(self, svcCtx: ServiceContext, payload, splitter: TextSplitter):
        super().__init__(task_type=PDFTaskType.PARSE)
        self.svcCtx = svcCtx
        self._payload = payload
        self._splitter = splitter
        
    def handle(self):
        paper_meta = None
        try:
            start = time.time()
            
            paper_meta = KqPaperMeta.from_json(self._payload)
            if paper_meta.user_id == 0 or not paper_meta.id:
                raise ValueError("invalid paper_meta, id or user_id invalid")
            if paper_meta.parse_status == ParseStatus.SUCCESS:
                raise RuntimeError("paper file already parse")
            
            self._update_parse_status(paper_meta, ParseStatus.PROCESSING)
        
            with self._get_paper_file(paper_meta) as tmp_file:
                md_text = self._parse_pdf2md(tmp_file)
                
                segments = []
                for text_segment in self._splitter.split(md_text, {"user_id": paper_meta.user_id, "paper_file_id": paper_meta.id}):
                    segments.append(text_segment)
                    
                    if len(segments) >= 10:  # 每10条处理一次
                        self._process_segments(paper_meta, segments)
                        segments = []
                # 处理剩余片段
                self._process_segments(paper_meta, segments)
            
            self._update_parse_status(paper_meta, ParseStatus.SUCCESS)
            duration_ms = round((time.time() - start) * 1000, 1)
            logger.my_logger().info(f'【PDFParseTask】parse file id: {paper_meta.id} success', extra={'duration': duration_ms})
        except Exception as e:
            if paper_meta is not None and not (paper_meta.user_id == 0 or not paper_meta.id):
                logger.my_logger().info(f'【PDFParseTask】parse file id: {paper_meta.id} failed, Exception {str(e)} happen')
                self._update_parse_status(paper_meta, ParseStatus.FAILED)
            else:
                logger.my_logger().info(f'【PDFParseTask】parse file failed, Exception {str(e)} happen')
                
    def _process_segments(self, paper_meta: KqPaperMeta, segments: List[TextSegment]):
        """处理并存储文本片段"""
        if not segments:
            return
        
        texts = [s.text for s in segments]
        vecs = self.svcCtx.embeddings.embed_documents(texts) # todo: 后续多分割器流程中，考虑将编码工作放到segment里面
        insert_datas = [
            self.svcCtx.paper_segment_model.paper_segment_data(
                user_id=paper_meta.user_id,
                paper_file_id=paper_meta.id,
                paper_file_name=paper_meta.file_name,
                text=s.text,
                embedding_vector=vecs[i],
                splitter_id=s.splitter_id,
                meta=s.meta
            ) for i, s in enumerate(segments)
        ]
        self.svcCtx.paper_segment_model.insert_many(insert_datas)
            
    def _update_parse_status(self, paper_meta: KqPaperMeta, status: ParseStatus):
        self.svcCtx.paperfile_stub.UpdateParseStatus(paperfile_pb2.UpdateParseStatusReq(
                Id=paper_meta.id,
                UserId=paper_meta.user_id,
                VectorDBInfo=paperfile_pb2.VectorDBInfo(ParseStatus=status.value)
            )
        )
    
    @contextmanager
    def _get_paper_file(self, paper_meta: KqPaperMeta):
        tmp_file = None
        try:
            tmp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=True)
            
            # 下载文件并强制同步写入
            self.svcCtx.minio_client.fget_object(
                paper_meta.minio_bucket,
                paper_meta.minio_object,
                tmp_file.name
            )
            tmp_file.flush()
            os.fsync(tmp_file.fileno())  # 确保物理写入，据说怕因为缓冲的原因导致后续PyMuPdf处理读文件读不全
            
            # 返回前重置文件指针
            tmp_file.seek(0)
            yield tmp_file
            
        except S3Error as e:
            error_msg = f"【PDFParseTask】_get_paper_file: MinIO操作失败: {e.code} (bucket={paper_meta.minio_bucket} object={paper_meta.minio_object})"
            logger.my_logger().error(error_msg)
        except Exception as e:
            error_msg = f"【PDFParseTask】_get_paper_file: 异常: {type(e).__name__} (原始错误: {str(e)})"
            raise RuntimeError(error_msg) from e
        finally:
            # 确保生成器结束后关闭文件
            if tmp_file and not tmp_file.closed:
                tmp_file.close()
                
    def _parse_pdf2md(self, tmp_file) -> str:
        doc = fitz.open(tmp_file.name)
        md_pages = []

        for page in doc:
            page_width = page.rect.width
            page_height = page.rect.height
            text_dict = page.get_text("dict")
            
            # 初始化裁剪高度
            crop_height = page_height
            
            # 检测页码（仅检查最后一个block的最后一行）
            if text_dict["blocks"]:
                last_block = text_dict["blocks"][-1]
                if "lines" in last_block and last_block["lines"]:
                    last_line = last_block["lines"][-1]
                    if "spans" in last_line and last_line["spans"]:
                        last_span = last_line["spans"][-1]
                        if self._is_page_number(last_span, page_height, page_width):
                            crop_height = last_span["bbox"][1] - (last_span["size"] * 0.5)

            # 创建裁剪页面
            tmp_doc = fitz.open()
            new_page = tmp_doc.new_page(
                width=page.rect.width,
                height=crop_height
            )
            new_page.show_pdf_page(
                new_page.rect,
                doc,
                page.number,
                clip=fitz.Rect(0, 0, page.rect.width, crop_height)
            )
            
            md_pages.append(pymupdf4llm.to_markdown(tmp_doc))
            
        return "\n".join(md_pages)
    
    def _is_page_number(self, span, page_height, page_width) -> bool:
        """页码判断逻辑"""
        # 预处理：去除首尾空白字符（包括换行符）
        text = span["text"].strip()
        
        # 规则1：位置检测（底部10%区域）
        is_bottom = span["bbox"][3] > page_height * 0.9
        
        # 规则2：宽度检测（不超过页面宽度20%）
        is_narrow = (span["bbox"][2] - span["bbox"][0]) < page_width * 0.2
        
        # 增强版规则3：内容检测（支持多种页码格式）
        # 匹配模式包括：
        # - 纯数字 "1", "23"
        # - "Page 1", "page 5" （大小写不敏感）
        # - 可能包含末尾空格/换行（已被strip处理）
        page_pattern = r'^(?:page\s*)?\d+\s*$'
        is_valid_format = bool(re.match(page_pattern, text, re.IGNORECASE))
        
        return is_bottom and is_narrow and is_valid_format
